#!/usr/bin/env python3
"""
FastAPI接口测试脚本
用于测试BERT意图识别FastAPI服务的所有端点
"""

import requests
import json
import time
import sys
from typing import Dict, List, Any
from colorama import init, Fore, Style
import argparse

# 初始化colorama
init()

class FastAPITester:
    """FastAPI接口测试器"""
    
    def __init__(self, base_url: str = "http://localhost:5000"):
        self.base_url = base_url.rstrip('/')
        self.session = requests.Session()
        self.test_results = []
        
    def log_info(self, message: str):
        """输出信息日志"""
        print(f"{Fore.GREEN}[INFO]{Style.RESET_ALL} {message}")
        
    def log_error(self, message: str):
        """输出错误日志"""
        print(f"{Fore.RED}[ERROR]{Style.RESET_ALL} {message}")
        
    def log_warning(self, message: str):
        """输出警告日志"""
        print(f"{Fore.YELLOW}[WARNING]{Style.RESET_ALL} {message}")
        
    def log_success(self, message: str):
        """输出成功日志"""
        print(f"{Fore.CYAN}[SUCCESS]{Style.RESET_ALL} {message}")
        
    def make_request(self, method: str, endpoint: str, **kwargs) -> Dict[str, Any]:
        """发送HTTP请求"""
        url = f"{self.base_url}{endpoint}"
        
        try:
            start_time = time.time()
            response = self.session.request(method, url, **kwargs)
            response_time = time.time() - start_time
            
            # 尝试解析JSON响应
            try:
                data = response.json()
            except json.JSONDecodeError:
                data = {"error": "Invalid JSON response", "text": response.text}
            
            result = {
                "endpoint": endpoint,
                "method": method,
                "status_code": response.status_code,
                "response_time": response_time,
                "data": data,
                "success": 200 <= response.status_code < 300
            }
            
            self.test_results.append(result)
            return result
            
        except requests.exceptions.ConnectionError:
            self.log_error(f"无法连接到 {url}")
            return {
                "endpoint": endpoint,
                "method": method,
                "status_code": 0,
                "response_time": 0,
                "data": {"error": "Connection failed"},
                "success": False
            }
        except Exception as e:
            self.log_error(f"请求失败: {e}")
            return {
                "endpoint": endpoint,
                "method": method,
                "status_code": 0,
                "response_time": 0,
                "data": {"error": str(e)},
                "success": False
            }
    
    def test_health_check(self):
        """测试健康检查端点"""
        self.log_info("测试健康检查...")
        
        result = self.make_request("GET", "/")
        
        if result["success"]:
            self.log_success(f"健康检查通过 (HTTP {result['status_code']})")
            self.log_info(f"响应时间: {result['response_time']:.3f}s")
            if "status" in result["data"]:
                self.log_info(f"服务状态: {result['data']['status']}")
        else:
            self.log_error(f"健康检查失败 (HTTP {result['status_code']})")
            
        return result["success"]
    
    def test_detailed_health(self):
        """测试详细健康检查"""
        self.log_info("测试详细健康检查...")
        
        result = self.make_request("GET", "/health/detailed")
        
        if result["success"]:
            self.log_success("详细健康检查通过")
            data = result["data"]
            if "model_loaded" in data:
                status = "已加载" if data["model_loaded"] else "未加载"
                self.log_info(f"模型状态: {status}")
        else:
            self.log_error("详细健康检查失败")
            
        return result["success"]
    
    def test_api_docs(self):
        """测试API文档端点"""
        self.log_info("测试API文档端点...")
        
        # 测试Swagger UI
        swagger_result = self.make_request("GET", "/docs")
        if swagger_result["success"]:
            self.log_success("Swagger UI 可访问")
        else:
            self.log_warning("Swagger UI 不可访问")
        
        # 测试ReDoc
        redoc_result = self.make_request("GET", "/redoc")
        if redoc_result["success"]:
            self.log_success("ReDoc 可访问")
        else:
            self.log_warning("ReDoc 不可访问")
        
        # 测试OpenAPI规范
        openapi_result = self.make_request("GET", "/openapi.json")
        if openapi_result["success"]:
            self.log_success("OpenAPI规范可访问")
            if "openapi" in openapi_result["data"]:
                version = openapi_result["data"].get("openapi", "unknown")
                self.log_info(f"OpenAPI版本: {version}")
        else:
            self.log_warning("OpenAPI规范不可访问")
        
        return swagger_result["success"] and redoc_result["success"]
    
    def test_model_info(self):
        """测试模型信息端点"""
        self.log_info("测试模型信息...")
        
        result = self.make_request("GET", "/model/info")
        
        if result["success"]:
            self.log_success("模型信息获取成功")
            data = result["data"]
            if "model_info" in data:
                model_info = data["model_info"]
                if "supported_intents" in model_info:
                    intents = model_info["supported_intents"]
                    self.log_info(f"支持的意图数量: {len(intents)}")
                    self.log_info(f"意图列表: {', '.join(intents)}")
        else:
            self.log_error("模型信息获取失败")
            
        return result["success"]
    
    def test_single_prediction(self):
        """测试单个预测"""
        self.log_info("测试单个文本预测...")
        
        test_cases = [
            {"text": "你好", "expected_intent": "问候"},
            {"text": "明天天气怎么样", "expected_intent": "查询天气"},
            {"text": "我要订票", "expected_intent": "订票"},
            {"text": "播放音乐", "expected_intent": "播放音乐"},
            {"text": "再见", "expected_intent": "再见"}
        ]
        
        success_count = 0
        
        for i, case in enumerate(test_cases, 1):
            self.log_info(f"测试用例 {i}: '{case['text']}'")
            
            payload = {
                "text": case["text"],
                "return_confidence": True,
                "return_all_scores": True
            }
            
            result = self.make_request("POST", "/predict", json=payload)
            
            if result["success"]:
                data = result["data"]
                predicted_intent = data.get("intent", "unknown")
                confidence = data.get("confidence", 0)
                
                self.log_success(f"预测成功: {predicted_intent} (置信度: {confidence:.3f})")
                
                if predicted_intent == case["expected_intent"]:
                    self.log_success("预测结果符合预期")
                    success_count += 1
                else:
                    self.log_warning(f"预测结果不符合预期，期望: {case['expected_intent']}")
                    
                # 显示所有分数
                if "all_scores" in data:
                    self.log_info("所有意图分数:")
                    for intent, score in data["all_scores"].items():
                        self.log_info(f"  {intent}: {score:.3f}")
                        
            else:
                self.log_error(f"预测失败: {result['data']}")
        
        self.log_info(f"单个预测测试完成: {success_count}/{len(test_cases)} 符合预期")
        return success_count == len(test_cases)
    
    def test_batch_prediction(self):
        """测试批量预测"""
        self.log_info("测试批量预测...")
        
        test_texts = [
            "早上好",
            "今天会下雨吗",
            "我想听歌",
            "帮我订张票",
            "再见了"
        ]
        
        payload = {
            "texts": test_texts,
            "batch_size": 3
        }
        
        result = self.make_request("POST", "/predict/batch", json=payload)
        
        if result["success"]:
            data = result["data"]
            results = data.get("results", [])
            processing_time = data.get("processing_time", 0)
            
            self.log_success(f"批量预测成功: {len(results)} 个结果")
            self.log_info(f"处理时间: {processing_time:.3f}s")
            self.log_info(f"平均每个文本: {processing_time/len(test_texts):.3f}s")
            
            for i, res in enumerate(results):
                text = res.get("text", "")
                intent = res.get("intent", "")
                self.log_info(f"  {i+1}. '{text}' -> {intent}")
                
        else:
            self.log_error(f"批量预测失败: {result['data']}")
            
        return result["success"]
    
    def test_conversation_prediction(self):
        """测试对话预测"""
        self.log_info("测试对话序列预测...")
        
        conversation = [
            "你好",
            "今天天气怎么样",
            "那明天呢",
            "好的，谢谢",
            "再见"
        ]
        
        payload = {
            "conversation": conversation
        }
        
        result = self.make_request("POST", "/predict/conversation", json=payload)
        
        if result["success"]:
            data = result["data"]
            conversation_results = data.get("conversation_results", [])
            
            self.log_success(f"对话预测成功: {len(conversation_results)} 轮")
            
            for i, turn_result in enumerate(conversation_results):
                text = turn_result.get("text", "")
                intent = turn_result.get("intent", "")
                self.log_info(f"  第{i+1}轮: '{text}' -> {intent}")
                
        else:
            self.log_error(f"对话预测失败: {result['data']}")
            
        return result["success"]
    
    def test_distribution_analysis(self):
        """测试意图分布分析"""
        self.log_info("测试意图分布分析...")
        
        test_texts = [
            "你好", "早上好", "下午好",  # 问候
            "明天天气如何", "今天会下雨吗",  # 查询天气
            "我要订票", "帮我买张机票",  # 订票
            "播放音乐", "来首歌",  # 播放音乐
            "再见", "拜拜"  # 再见
        ]
        
        payload = {
            "texts": test_texts
        }
        
        result = self.make_request("POST", "/analyze/distribution", json=payload)
        
        if result["success"]:
            data = result["data"]
            distribution = data.get("distribution", {})
            
            self.log_success("意图分布分析成功")
            
            if "intent_counts" in distribution:
                self.log_info("意图分布:")
                for intent, count in distribution["intent_counts"].items():
                    percentage = distribution.get("intent_percentages", {}).get(intent, 0)
                    self.log_info(f"  {intent}: {count} ({percentage:.1f}%)")
                    
        else:
            self.log_error(f"意图分布分析失败: {result['data']}")
            
        return result["success"]
    
    def test_benchmark(self):
        """测试性能基准"""
        self.log_info("测试性能基准...")
        
        test_texts = [
            "你好，我想查询天气",
            "帮我订张票",
            "播放一首歌"
        ]
        
        payload = {
            "test_texts": test_texts,
            "num_runs": 3
        }
        
        result = self.make_request("POST", "/benchmark", json=payload)
        
        if result["success"]:
            data = result["data"]
            benchmark_stats = data.get("benchmark_stats", {})
            
            self.log_success("性能基准测试成功")
            
            if "average_time" in benchmark_stats:
                avg_time = benchmark_stats["average_time"]
                min_time = benchmark_stats.get("min_time", 0)
                max_time = benchmark_stats.get("max_time", 0)
                
                self.log_info(f"平均响应时间: {avg_time:.3f}s")
                self.log_info(f"最快响应时间: {min_time:.3f}s")
                self.log_info(f"最慢响应时间: {max_time:.3f}s")
                
        else:
            self.log_error(f"性能基准测试失败: {result['data']}")
            
        return result["success"]
    
    def test_input_validation(self):
        """测试输入验证"""
        self.log_info("测试输入验证...")
        
        # 测试空文本
        result = self.make_request("POST", "/predict", json={"text": ""})
        if result["status_code"] == 422:  # FastAPI validation error
            self.log_success("空文本验证正确")
        else:
            self.log_warning("空文本验证未按预期工作")
        
        # 测试过长文本
        long_text = "这是一个非常长的文本" * 100
        result = self.make_request("POST", "/predict", json={"text": long_text})
        if result["status_code"] == 422:
            self.log_success("过长文本验证正确")
        else:
            self.log_warning("过长文本验证未按预期工作")
        
        # 测试无效JSON
        result = self.make_request("POST", "/predict", 
                                 data="invalid json", 
                                 headers={"Content-Type": "application/json"})
        if result["status_code"] >= 400:
            self.log_success("无效JSON验证正确")
        else:
            self.log_warning("无效JSON验证未按预期工作")
        
        return True
    
    def run_all_tests(self):
        """运行所有测试"""
        self.log_info("=" * 60)
        self.log_info("开始FastAPI接口完整测试")
        self.log_info("=" * 60)
        
        tests = [
            ("健康检查", self.test_health_check),
            ("详细健康检查", self.test_detailed_health),
            ("API文档", self.test_api_docs),
            ("模型信息", self.test_model_info),
            ("单个预测", self.test_single_prediction),
            ("批量预测", self.test_batch_prediction),
            ("对话预测", self.test_conversation_prediction),
            ("分布分析", self.test_distribution_analysis),
            ("性能基准", self.test_benchmark),
            ("输入验证", self.test_input_validation),
        ]
        
        results = {}
        total_start_time = time.time()
        
        for test_name, test_func in tests:
            self.log_info(f"\n--- {test_name} ---")
            try:
                results[test_name] = test_func()
            except Exception as e:
                self.log_error(f"{test_name}测试出现异常: {e}")
                results[test_name] = False
        
        total_time = time.time() - total_start_time
        
        # 生成测试报告
        self.generate_report(results, total_time)
    
    def generate_report(self, results: Dict[str, bool], total_time: float):
        """生成测试报告"""
        self.log_info("\n" + "=" * 60)
        self.log_info("测试报告")
        self.log_info("=" * 60)
        
        passed = sum(results.values())
        total = len(results)
        
        self.log_info(f"总测试数: {total}")
        self.log_success(f"通过: {passed}")
        if total - passed > 0:
            self.log_error(f"失败: {total - passed}")
        
        self.log_info(f"总耗时: {total_time:.2f}s")
        self.log_info(f"成功率: {passed/total*100:.1f}%")
        
        # 详细结果
        self.log_info("\n详细结果:")
        for test_name, success in results.items():
            status = "✅ PASS" if success else "❌ FAIL"
            print(f"  {test_name:20} {status}")
        
        # 性能统计
        if self.test_results:
            response_times = [r["response_time"] for r in self.test_results if r["success"]]
            if response_times:
                avg_time = sum(response_times) / len(response_times)
                max_time = max(response_times)
                min_time = min(response_times)
                
                self.log_info(f"\n性能统计:")
                self.log_info(f"  平均响应时间: {avg_time:.3f}s")
                self.log_info(f"  最快响应时间: {min_time:.3f}s")
                self.log_info(f"  最慢响应时间: {max_time:.3f}s")
        
        # 建议
        if passed == total:
            self.log_success("🎉 所有测试通过！API服务运行正常。")
        else:
            self.log_warning("⚠️ 部分测试失败，请检查服务状态和配置。")
        
        return passed == total


def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="FastAPI接口测试工具")
    parser.add_argument("--url", default="http://localhost:5000", 
                      help="API服务地址 (默认: http://localhost:5000)")
    parser.add_argument("--test", choices=[
        "health", "docs", "model", "predict", "batch", 
        "conversation", "distribution", "benchmark", "validation", "all"
    ], default="all", help="指定测试类型 (默认: all)")
    
    args = parser.parse_args()
    
    tester = FastAPITester(args.url)
    
    # 首先检查服务是否可访问
    if not tester.test_health_check():
        tester.log_error(f"无法连接到API服务: {args.url}")
        tester.log_info("请确保：")
        tester.log_info("1. API服务正在运行")
        tester.log_info("2. 服务地址正确")
        tester.log_info("3. 防火墙设置允许访问")
        sys.exit(1)
    
    # 运行指定测试
    if args.test == "all":
        success = tester.run_all_tests()
    else:
        test_map = {
            "health": tester.test_health_check,
            "docs": tester.test_api_docs,
            "model": tester.test_model_info,
            "predict": tester.test_single_prediction,
            "batch": tester.test_batch_prediction,
            "conversation": tester.test_conversation_prediction,
            "distribution": tester.test_distribution_analysis,
            "benchmark": tester.test_benchmark,
            "validation": tester.test_input_validation,
        }
        
        test_func = test_map.get(args.test)
        if test_func:
            success = test_func()
        else:
            tester.log_error(f"未知测试类型: {args.test}")
            sys.exit(1)
    
    sys.exit(0 if success else 1)


if __name__ == "__main__":
    main() 